import torch
import numpy as np
import pandas as pd

class TrainerReturnedObject:
    def __init__(self,
                 filenames: torch.Tensor = None,
                 targets: torch.Tensor = None,
                 preds: torch.Tensor = None,
                 losses: torch.Tensor = None) -> None:
        self.filenames = self.to_numpy(filenames)
        self.targets = self.to_numpy(targets)
        self.preds = self.to_numpy(preds)
        self.losses = 0.0 if losses is None else losses * len(filenames)
        self.num = 0 if filenames is None else len(filenames)

    @property
    def preds_argmax(self) -> np.ndarray:
        return np.argmax(self.preds, axis=1)

    def predict_table(self):
        df = pd.DataFrame()
        df['filename'] = list(self.filenames)
        df['predict'] = self.preds_argmax.tolist()
        df['true'] = self.targets.tolist()
        df['correct'] = (self.preds_argmax == self.targets).tolist()
        return df

    def to_numpy(self, tensor_obj: torch.Tensor):
        if isinstance(tensor_obj, torch.Tensor):
            return tensor_obj.cpu().detach().numpy()
        else:
            return tensor_obj

    def append(self, obj_dst: np.ndarray, obj_src: np.ndarray):
        if obj_dst is None:
            return obj_src
        else:
            return np.append(obj_dst, obj_src, axis=0)

    def update(self, trobj):
        one_batch_num = len(trobj.filenames)
        self.filenames = self.append(self.filenames,
                                     self.to_numpy(trobj.filenames))
        self.targets = self.append(self.targets, self.to_numpy(trobj.targets))
        self.preds = self.append(self.preds, self.to_numpy(trobj.preds))
        # self.losses = self.losses + trobj.losses * one_batch_num
        self.losses = self.losses + trobj.losses
        self.num += one_batch_num


class Trainer(object):
    """
    训练器
    """
    def __init__(self, model: torch.nn.Module, loss_func, optimizer, device):
        """
        Args:
            loss_func: 损失函数
            optimizer: 优化器
            model: 模型
            device: 数据运算的设备
        """
        self.loss_func = loss_func
        self.optimizer = optimizer
        self.model = model.to(device=device)
        self.device = device

    def _forward(self, samples, is_train=True):
        filenames, imgs, targets = samples
        #将数据移至device
        imgs = imgs.to(self.device)
        targets = targets.to(self.device)
        #前向传播
        preds = self.model(imgs)
        losses = self.loss_func(preds, targets)
        if is_train:
            #反向传播
            self.optimizer.zero_grad()  #梯度置0，否则下一轮的梯度会累积之前的梯度
            losses.backward()
            self.optimizer.step()  #优化器学习率更新
        return TrainerReturnedObject(filenames, targets, preds, losses)

    def train(self, samples):
        self.model.train()
        # with torch.enable_grad():
        #     return self._forward(samples, True)
        return self._forward(samples, True)

    def validate(self, samples):
        self.model.eval()
        with torch.no_grad():
            return self._forward(samples, False)